#!/usr/bin/env python

import cv2
import numpy as np
import os
import glob

# Defining the dimensions of checkerboard
CHECKERBOARD = (6,9)
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001) 

# Creating vector to store vectors of 3D points for each checkerboard image
objpoints = []
# Creating vector to store vectors of 2D points for each checkerboard image
imgpoints = [] 


# Defining the world coordinates for 3D points
objp = np.zeros((1, CHECKERBOARD[0]*CHECKERBOARD[1], 3), np.float32)
objp[0,:,:2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
prev_img_shape = None
print("type(objp): ", type(objp))

# 提取存储在给定目录中的单个图像的路径
validPicNum = 0
images = glob.glob('./images/*.jpg')
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    # 寻找棋盘角点。
    # 如果在棋盘中找到了所需的角数，则 ret = true
    ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD, cv2.CALIB_CB_ADAPTIVE_THRESH+
    	cv2.CALIB_CB_FAST_CHECK+cv2.CALIB_CB_NORMALIZE_IMAGE)
    
    """
    如果检测到所需的角点数量，我们优化像素坐标并把他们显示在棋盘上
    """
    if ret == True:
        # 收集空间3D点
        # print("objp"+str(validPicNum)+": ", objp)
        objpoints.append(objp)

        # 优化给定的 2d 像素坐标
        corners2 = cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        
        # 收集给定的 2d 像素
        # print("corners"+str(validPicNum)+": ", corners2)
        imgpoints.append(corners2)

        # 绘制并显示角点
        img = cv2.drawChessboardCorners(img, CHECKERBOARD, corners2,ret)
        
        # 保存图片
        # cv2.imwrite("img"+str(validPicNum)+".jpg", img)
        
        validPicNum += 1
    
    cv2.imshow('img',img)
    cv2.waitKey(33)

cv2.destroyAllWindows()

h,w = img.shape[:2]
print("\n image size(h, w): (", h, " ", w, ")")
print(h, " ", w)
print("ValidPicNum: ", validPicNum)

"""
传递已知的 3D 点以及与之对齐的像素坐标执行相机标定
"""
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1],None,None)

print("\n Camera matrix : ")
print(mtx)
print("\n dist : ")
print(dist)
print("\n rvecs : ")
print(rvecs)
print("\n tvecs : ")
print(tvecs, "; len: ", len(tvecs))
